Super resolution on medical images:

Idea of this project is to create a model that is able to output a high resolution image from a low resolution one, based on the assumption that it can be trained on both high and low resolution inputs. See github for report and more information.

Link to doc Fastai

Set-up:

Start with importing the right functions and modules.

Global imports:

Change simulation number :) this number will be used to save your data so change it when you run with new parameters

In [77]:
sim = 9
In [2]:
import pandas as pd
import fastai
from fastai.vision import *
from fastai.callbacks import *
from fastai.utils.mem import *
import tensorflow as tf
from torchvision.models import vgg16_bn
import torchvision.transforms as transforms
from skimage import measure
from numpy import linalg as LA
import shutil

# Additional personal functions necessary to run the notebook. 
from  add_functions.model_create_data_functions import *
from  add_functions.plot_fx_model import *
from  add_functions.evaluation_metrics import *

import warnings
warnings.filterwarnings('ignore')
In [3]:
%matplotlib inline
%load_ext autoreload
%autoreload 2

Set CUDA if available.

In [4]:
torch.cuda.set_device(3)
In [5]:
device = torch.device("cuda:3" if torch.cuda.is_available() else "cpu")
print(device)
cuda:3

Paths to data:

Input images:

  • HR images: (3, 502, 672) original (3, 1004, 1344) cut in 4 pieces

MR and LR will be created below.

  • LR images: (3, 125, 167) original (3, 500, 669) cut in 4 pieces, downsampled by 2.
  • MR images: (3, 250, 334) original (3, 500, 669) cut in 4 pieces

Paths to data:

In [6]:
path = Path('../../../../../SCRATCH2/marvande/data/train/HR/')

# path to HR data:
path_hr = path / 'HR_patches_train/tiff_files'

# path to save LR and MR images:
path_lr = path / 'small-125/train'
path_mr = path / 'small-250/train'

# path to test images:
path_test = path/'HR_patches_test/cut_images/tiff_files/'

# path to save LR and MR test images:
path_lr_test = path / 'small-125/test'
path_mr_test = path / 'small-250/test'

assert path.exists(), f"need dataset @ {path}"
assert path_hr.exists()

Run cell below if you want to resize the data each time you run the notebook. Otherwise comment.

In [7]:
! rm -r ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train
! rm -r ../../../../../SCRATCH2/marvande/data/train/HR/small-125/test
! rm -r ../../../../../SCRATCH2/marvande/data/train/HR/small-250/train
! rm -r ../../../../../SCRATCH2/marvande/data/train/HR/small-250/test

Create LR and MR data:

From this HR data, we create LR versions in path_lr and path_mr for training and testing. LR images will be of size (3, 125, 167) while MR images of size (3, 250, 334). Both will be saved as lower quality than HR.

Resizing training HR to MR, LR:

In [8]:
# create smaller image sets of lower quality the first time this nb is run:
new_size_lr = (167, 125)
sets = [(path_lr, new_size_lr), (path_mr, (334, 250))]
for p, size in sets:
    if not p.exists():
        print(f"resizing to {size} into {p} from {path_hr}")
        parallel(partial(resize_one, path=p, size=size, rel_to_path=path_hr),
                 ImageList.from_folder(path_hr).items)
resizing to (167, 125) into ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train from ../../../../../SCRATCH2/marvande/data/train/HR/HR_patches_train/tiff_files
resizing to (334, 250) into ../../../../../SCRATCH2/marvande/data/train/HR/small-250/train from ../../../../../SCRATCH2/marvande/data/train/HR/HR_patches_train/tiff_files

Resize test HR to MR, LR:

In [9]:
# create smaller tests sets the first time this nb is run:
new_size_lr = (167, 125)
sets = [(path_lr_test, new_size_lr), (path_mr_test, (334, 250))]
for p, size in sets:
    if not p.exists():
        print(f"resizing to {size} into {p} from {path_test}")
        parallel(partial(resize_one, path=p, size=size, rel_to_path=path_test),
                 ImageList.from_folder(path_test).items)
resizing to (167, 125) into ../../../../../SCRATCH2/marvande/data/train/HR/small-125/test from ../../../../../SCRATCH2/marvande/data/train/HR/HR_patches_test/cut_images/tiff_files
resizing to (334, 250) into ../../../../../SCRATCH2/marvande/data/train/HR/small-250/test from ../../../../../SCRATCH2/marvande/data/train/HR/HR_patches_test/cut_images/tiff_files

 Evaluate resizing in a test image:

To have an idea about what exactly happens to the pixels and images during resizing, we test it out on one image.

In [10]:
path_t = '../../../../../SCRATCH2/marvande/data/train/HR/HR_patches_train/tiff_files/'
file_p = path_t + '0013_[39667,16250]_part_1_1_.tif' 
file_name = 'HR_test.tif'
file = PIL.Image.open(file_p)
file.save('test/HR/'+file_name)

p_hr = Path('test/HR/')
p_lr = Path('test/LR')
p_mr = Path('test/MR')

size_lr_test = (167, 125)

# create smaller image sets of lower quality the first time this nb is run:
sets = [(Path('test/LR/'), size_lr_test), (Path('test/MR/'), (334, 250))]
for p, size in sets:
    if not p.exists():
        print(f"resizing to {size} into {p} from {p_hr}")
        parallel(partial(resize_one, path=p, size=size, rel_to_path = p_hr),
                 ImageList.from_folder('test/HR/').items)
In [11]:
# Compare images after resizing:
HR_tens = scaler(open_image(p_hr/file_name).data)

img_mr = PIL.Image.open(p_mr/file_name)
img_mr = img_mr.resize((672, 502), resample=PIL.Image.BICUBIC).convert('RGB')
MR_tens = scaler(trans1(img_mr))

#evaluate norm between MR, HR:
evaluate_metrics(HR_tens.numpy(), MR_tens.numpy(), 'MR VS HR')
print('')

img_lr = PIL.Image.open(p_lr/file_name)
img_lr = img_lr.resize((672, 502), resample=PIL.Image.BICUBIC).convert('RGB')
LR_tens = scaler(trans1(img_lr))

#evaluate norm between MR, HR:
evaluate_metrics(HR_tens.numpy(), LR_tens.numpy(), 'LR VS HR')
print('')

plot_triple(MR_tens.numpy()[0,:,:], HR_tens.numpy()[0,:,:], LR_tens.numpy()[0,:,:], 'MR','HR', 'ups LR')
MR VS HR: Frobenius diff: 9.8090 frob, Nuclear diff: 165.2214
MR VS HR: MAE: 3.195315, MSE: 19.6178, NMSE: 0.0193, PSNR: 35.2043, SSIM: 0.9872

LR VS HR: Frobenius diff: 21.7091 frob, Nuclear diff: 310.5509
LR VS HR: MAE: 5.484263, MSE: 92.4003, NMSE: 0.0419, PSNR: 28.4741, SSIM: 0.9428

In [12]:
plot_triple(LR_tens.numpy()[0,0:10,0:10], MR_tens.numpy()[0,0:10,0:10],
            HR_tens.numpy()[0,0:10,0:10], 'LR', 'MR', 'HR')

Training, validation and test data:

Data phase 1:

During phase 1, the model is trained on images of size (3, 250, 334). Therefore HR images (originally of size (3, 502, 672)) are down-sampled by approximately a factor 2 and LR images (originally of size (3, 125, 334)) up-sampled by 2.

In [13]:
# first phase data is LR data upsampled by factor 2:
bs, size = 15, (250, 334)
data_phase1 = get_data(bs, size, path_lr_test, path_lr, path_hr)
print(f'Test origin {path_lr_test}')
data_phase1
Test origin ../../../../../SCRATCH2/marvande/data/train/HR/small-125/test
Out[13]:
ImageDataBunch;

Train: LabelList (1851 items)
x: ImageImageList
Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334)
y: ImageList
Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334)
Path: ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train;

Valid: LabelList (205 items)
x: ImageImageList
Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334)
y: ImageList
Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334)
Path: ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train;

Test: LabelList (336 items)
x: ImageImageList
Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334),Image (3, 250, 334)
y: EmptyLabelList
,,,,
Path: ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train

Have a look at a few data from the validation data. Left are the LR while right the corresponding Hr images. Furthermore, we evaluate some metrics to look at their similarity.

In [14]:
LR_tens = data_phase1.one_batch(ds_type=DatasetType.Valid)[0][1]
HR_tens = data_phase1.one_batch(ds_type=DatasetType.Valid)[1][1]
print('LR data shape {} and HR shape {}'.format(list(LR_tens.numpy().shape),
                                                list(HR_tens.numpy().shape)))
evaluate_metrics(LR_tens.numpy(), HR_tens.numpy(), 'LR vs HR phase 1')
plot_double(scaler(LR_tens).numpy()[0, :, :], scaler(HR_tens).numpy()[0, :, :], 'LR', 'HR')
LR data shape [3, 250, 334] and HR shape [3, 250, 334]
LR vs HR phase 1: Frobenius diff: 5.7585 frob, Nuclear diff: 71.1713
LR vs HR phase 1: MAE: 3.841993, MSE: 26.1032, NMSE: 0.0229, PSNR: 33.9639, SSIM: 0.9651

Data phase 2:

During phase 2, the model is trained on images of size (3, 502, 672). Therefore LR images (originally of size (3, 125, 334)) up-sampled by approximately 4.

In [15]:
bs, new_size = 4, (502, 672)
data_phase2 = get_data(bs, new_size, path_lr_test, path_lr, path_hr)
data_phase2
Out[15]:
ImageDataBunch;

Train: LabelList (1851 items)
x: ImageImageList
Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672)
y: ImageList
Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672)
Path: ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train;

Valid: LabelList (205 items)
x: ImageImageList
Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672)
y: ImageList
Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672)
Path: ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train;

Test: LabelList (336 items)
x: ImageImageList
Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672),Image (3, 502, 672)
y: EmptyLabelList
,,,,
Path: ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train

Have a look at a few data from the validation data. Left are the LR while right the corresponding Hr images. Furthermore, we evaluate some metrics to look at their similarity.

In [16]:
LR_tens = data_phase2.one_batch(ds_type=DatasetType.Valid)[0][1]
HR_tens = data_phase2.one_batch(ds_type=DatasetType.Valid)[1][1]
print('LR data shape {} and HR shape {}'.format(list(LR_tens.numpy().shape),
                                                list(HR_tens.numpy().shape)))
evaluate_metrics(LR_tens.numpy(), HR_tens.numpy(), 'LR vs HR phase 1')
plot_double(
    scaler(LR_tens).numpy()[0, :, :],
    scaler(HR_tens).numpy()[0, :, :], 'LR', 'HR')
LR data shape [3, 502, 672] and HR shape [3, 502, 672]
LR vs HR phase 1: Frobenius diff: 17.4040 frob, Nuclear diff: 301.2155
LR vs HR phase 1: MAE: 7.807130, MSE: 62.9437, NMSE: 0.0355, PSNR: 30.1413, SSIM: 0.8970

Feature loss:

Create loss metrics used in the model.

In [17]:
def gram_matrix(x):
    """
    Gram matrix of a set of vectors in an inner 
    product space is the Hermitian matrix of inner products
    """
    n,c,h,w = x.size()
    x = x.view(n, c, -1)
    return (x @ x.transpose(1,2))/(c*h*w)

Select the data of an image and use it to compute its Gram matrix.

In [18]:
im = data_phase1.valid_ds[0][1]
t = im.data
t = torch.stack([t,t])
In [19]:
gram_matrix(t)
Out[19]:
tensor([[[0.2957, 0.3039, 0.2875],
         [0.3039, 0.3125, 0.2959],
         [0.2875, 0.2959, 0.2840]],

        [[0.2957, 0.3039, 0.2875],
         [0.3039, 0.3125, 0.2959],
         [0.2875, 0.2959, 0.2840]]])

We define a base loss as the L1 loss (F.l1_loss).

In [20]:
base_loss = F.l1_loss

Construct a pre-trained vgg16 model (Very Deep Convolutional Networks for Large-Scale Image Recognition). VGG is another Convolutional Neural Network (CNN) architecture devised in 2014, the 16 layer version is utilised in the loss function for training this model. VGG model. a network pretrained on ImageNet, is used to evaluate the generator model’s loss.

Further, we set requires_grad to False as this is useful when you want to freeze part of your model, or you know in advance that you are not going to use gradients w.r.t. some parameters.

The head of the VGG model is the final layers shown as fully connected and softmax in the above diagram. This head is ignored and the loss function uses the intermediate activations in the backbone of the network, which represent the feature detections.

Those activations can be found by looking through the VGG model to find all the max pooling layers. These are where the grid size changes and features are detected. So we need to select those layers.

Create a VGG model - just using the pre-trained model. In VGG, there's a attribute called .features which contains the convolutional part of the model. So vgg16_bn(True).features is the convolutional part of the VGG model. Because we don't need the head. We only want the intermediate activations. Then we'll check that on the GPU, we'll put it into eval mode because we're not training it. And we'll turn off requires_grad because we don't want to update the weights of this model. We're just using it for inference (i.e. for the loss).

In [21]:
vgg_m = vgg16_bn(True).features.cuda().eval()
requires_grad(vgg_m, False)

Select the layer IDs of MaxPool2d blocks. Enumerate through all the children of that model and find all of the max pooling layers, because in the VGG model that's where the grid size changes.

In [22]:
blocks = [i-1 for i,o in enumerate(children(vgg_m)) if isinstance(o,nn.MaxPool2d)]
blocks, [vgg_m[i] for i in blocks]
Out[22]:
([5, 12, 22, 32, 42],
 [ReLU(inplace=True),
  ReLU(inplace=True),
  ReLU(inplace=True),
  ReLU(inplace=True),
  ReLU(inplace=True)])

So we grab layer i-1. That's the layer before it changes. So there's our list of layer numbers just before the max pooling layers ([5, 12, 22, 32, 42]). All of those are ReLU's, not surprisingly. Those are where we want to grab some features from, and so we put that in blocks - it's just a list of ID's.

Create the feature loss from the model and layer ids selected above.

Source:

Main points:

  • Loss functions using these techniques can be used during the training of U-Net based model architectures
  • Feature loss: loss function used is similar to the loss function in the the paper, using VGG-16 but also combined with pixel mean squared error loss loss and gram matrix style loss
  • The training of a model can use this loss function based on the VGG model’s activations. The loss function remains fixed throughout the training unlike the critic part of a GAN
  • Feature loss: Feature map has 256 channels by 28 by 28. The activations at the same layer for the (target) original image and the generated image are compared using mean squared error or the least absolute error (L1) error for the base loss. These are feature losses. This error function uses L1 error. This allows the loss function to know what features are in the target ground truth image and to evaluate how well the model’s prediction’s features match these rather than only comparing pixel difference. This allows the model being trained with this loss function to produce much finer detail in the generated/predicted features and output.
  • Gram matrix style loss: A gram matrix defines a style with respect to specific content. By calculating the gram matrix for each feature activation in the target/ground truth image, it allows the style of that feature to be defined. If the same gram matrix is calculated from the activations of the predictions, the two can be compared to calculate how close the style of the feature prediction is to the target/ground truth image. A gram matrix is the matrix multiplication of the each of the activations and the activation matrix’s transpose. This enables the model to learn and generate predictions of images whose features look correct in their style and in context, with the end result looking more convincing and appear closer or the same as the target/ground truth.

Predictions from models trained with this loss function: The generated predictions from trained models using loss functions based on these techniques have both convincing fine detail and style. That style and fine detail may be different aspects of image quality be predicting fine pixel detail or the predicting correct colours.

In [23]:
class FeatureLoss(nn.Module):
    def __init__(self, m_feat, layer_ids, layer_wgts):
        super().__init__()
        self.m_feat = m_feat
        self.loss_features = [self.m_feat[i] for i in layer_ids]
        self.hooks = hook_outputs(self.loss_features, detach=False)
        self.wgts = layer_wgts
        self.metric_names = ['pixel',] + [f'feat_{i}' for i in range(len(layer_ids))
              ] + [f'gram_{i}' for i in range(len(layer_ids))]

    def make_features(self, x, clone=False):
        self.m_feat(x)
        return [(o.clone() if clone else o) for o in self.hooks.stored]
    
    def forward(self, input, target):
        out_feat = self.make_features(target, clone=True)
        in_feat = self.make_features(input)
        self.feat_losses = [base_loss(input,target)]
        #feat losses
        self.feat_losses += [base_loss(f_in, f_out)*w
                             for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
        #gram: 
        self.feat_losses += [base_loss(gram_matrix(f_in), gram_matrix(f_out))*w**2 * 5e3
                             for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)]
        self.metrics = dict(zip(self.metric_names, self.feat_losses))
        return sum(self.feat_losses) 
    
    def __del__(self): self.hooks.remove()
In [24]:
feat_loss = FeatureLoss(vgg_m, blocks[2:5], [5,15,2])

Training:

Set up:

In [25]:
# learning rate:
lr = 1e-3
#weight decay:
wd = 1e-3

def do_fit(save_name, lrs=slice(lr), pct_start=0.9):
    """
    do_fit: fits during 10 epochs with feature loss. 
    """
    learn.fit_one_cycle(10, lrs, pct_start=pct_start)
    learn.save(save_name, return_path=True)
    learn.show_results(rows=1, imgsize=10)
    learn.recorder.plot_losses()
    learn.recorder.plot_metrics()

In case pytorch is using too much memory and need to reset:

In [26]:
# delete all tensors and free cache:
for obj in gc.get_objects():
    if torch.is_tensor(obj):
        del obj
torch.cuda.empty_cache()
gc.collect()
#learn.destroy()

#get free memory (in MBs) for the currently selected gpu id, after emptying the cache
print(
    'free memory (in MBs) for the currently selected gpu id, after emptying the cache: ',
    gpu_mem_get_free_no_cache())

print(
    'used memory (in MBs) for the currently selected gpu id, after emptying the cache:',
    gpu_mem_get_used_no_cache())

gpu_mem_get_all()
free memory (in MBs) for the currently selected gpu id, after emptying the cache:  11026
used memory (in MBs) for the currently selected gpu id, after emptying the cache: 1186
Out[26]:
[GPUMemory(total=12212, free=11525, used=687),
 GPUMemory(total=12196, free=12185, used=10),
 GPUMemory(total=12196, free=12185, used=10),
 GPUMemory(total=12212, free=11026, used=1186)]
In [27]:
#resnet 34 given as an encoder in downsampling section of unet
arch = models.resnet34
learn = unet_learner(data_phase1,
                     arch,
                     wd=wd,
                     loss_func=feat_loss,
                     callback_fns=LossMetrics,
                     blur=True,
                     norm_type=NormType.Weight)
gc.collect();

Plots loss according to learning rate to pick the best learning rate.

In [28]:
learn.lr_find()
learn.recorder.plot()
0.00% [0/1 00:00<00:00]
epoch train_loss valid_loss pixel feat_0 feat_1 feat_2 gram_0 gram_1 gram_2 time

75.61% [93/123 01:57<00:37 3.2741]
LR Finder is complete, type {learner_name}.recorder.plot() to see the graph.

From this plot a learning rate of 1e-3 is chosen as it's approximately the middle of the descending curve.

In [29]:
lr = 1e-3

Phase 1:

During phase 1, the model is trained on images of size (3, 250, 334). Therefore HR images (originally of size (3, 502, 672)) are down-sampled by approximately a factor 2 and LR images (originally of size (3, 125, 334)) up-sampled by 2.

Phase 1 a:

As per usual, because we're using a pre-trained network in our U-Net, we start with frozen layers for the downsampling path, train for a while. As you can see, we get not only the loss, but also the pixel loss and the loss at each of our feature layers, and then also something we'll learn about in part 2 called Gram loss which I don't think anybody's used for super resolution before as far as I know.

lr = 0.01

In [31]:
do_fit('sim{}_1a'.format(sim), slice(lr * 10))
epoch train_loss valid_loss pixel feat_0 feat_1 feat_2 gram_0 gram_1 gram_2 time
0 1.062768 0.801133 0.096334 0.135392 0.134418 0.028890 0.208248 0.184166 0.013685 02:40
1 0.777333 0.677810 0.082140 0.117737 0.117370 0.025078 0.172102 0.151302 0.012081 02:38
2 0.740963 0.827157 0.153334 0.121644 0.122693 0.026486 0.213189 0.177257 0.012555 02:38
3 0.754527 0.738820 0.131806 0.119349 0.118345 0.025558 0.181891 0.149999 0.011873 02:38
4 0.747166 0.658180 0.092650 0.112167 0.111461 0.023826 0.166176 0.140361 0.011542 02:38
5 0.738996 0.693234 0.118644 0.114922 0.113359 0.025350 0.162196 0.146860 0.011903 02:38
6 0.758328 0.657018 0.098843 0.112345 0.111403 0.024148 0.160655 0.138091 0.011533 02:38
7 0.711967 0.628073 0.083350 0.109078 0.108469 0.023367 0.161588 0.131106 0.011116 02:38
8 0.670704 0.590957 0.086998 0.106876 0.105430 0.022919 0.137092 0.120898 0.010745 02:38
9 0.591649 0.525496 0.061218 0.104050 0.102213 0.022208 0.114493 0.111036 0.010277 02:38

Copy the learning object to another folder in order not to delete it when deleting paht_lr the data folder before resizing at the beginning of the notebook.

In [32]:
shutil.copy(
    path_lr/'models/sim{}_1a.pth'.format(sim),
    '../../../../../SCRATCH2/marvande/data/train/HR/models/')
Out[32]:
'../../../../../SCRATCH2/marvande/data/train/HR/models/sim9_1a.pth'

Phase 1 b:

We unfreeze and train some more. With smaller learning rates [1e-05 to 0.001]

In [33]:
learn.unfreeze()
In [34]:
do_fit('sim{}_1b'.format(sim), slice(1e-5,lr))
epoch train_loss valid_loss pixel feat_0 feat_1 feat_2 gram_0 gram_1 gram_2 time
0 0.567655 0.523721 0.060637 0.103519 0.101500 0.022076 0.115581 0.110149 0.010260 02:43
1 0.562572 0.520663 0.060556 0.103324 0.101331 0.022058 0.113894 0.109266 0.010235 02:43
2 0.560259 0.518266 0.061522 0.102890 0.101080 0.022077 0.112235 0.108228 0.010234 02:43
3 0.562876 0.513461 0.060338 0.102677 0.100802 0.022017 0.109812 0.107638 0.010177 02:43
4 0.557097 0.518190 0.059430 0.102452 0.100403 0.021859 0.114874 0.109039 0.010133 02:43
5 0.555372 0.507107 0.060039 0.102028 0.100321 0.021814 0.106585 0.106267 0.010054 02:43
6 0.549047 0.504988 0.059174 0.101165 0.099251 0.021580 0.107897 0.105965 0.009956 02:43
7 0.550850 0.506078 0.059697 0.100775 0.099135 0.021702 0.108698 0.106007 0.010064 02:44
8 0.548958 0.506370 0.060701 0.101193 0.099507 0.021734 0.106418 0.106786 0.010030 02:43
9 0.541038 0.494553 0.058767 0.100239 0.098628 0.021450 0.102248 0.103354 0.009867 02:43
In [35]:
shutil.copy(
    path_lr/'models/sim{}_1b.pth'.format(sim),
    '../../../../../SCRATCH2/marvande/data/train/HR/models/')
Out[35]:
'../../../../../SCRATCH2/marvande/data/train/HR/models/sim9_1b.pth'

Phase 2:

During phase 2, the model is trained on images of size (3, 502, 672). Therefore LR images (originally of size (3, 125, 334)) up-sampled by approximately 4.

In [27]:
torch.cuda.empty_cache()
gc.collect()
#learn.destroy()

#get free memory (in MBs) for the currently selected gpu id, after emptying the cache
print(
    'free memory (in MBs) for the currently selected gpu id, after emptying the cache: ',
    gpu_mem_get_free_no_cache())

print(
    'used memory (in MBs) for the currently selected gpu id, after emptying the cache:',
    gpu_mem_get_used_no_cache())

gpu_mem_get_all()
free memory (in MBs) for the currently selected gpu id, after emptying the cache:  11629
used memory (in MBs) for the currently selected gpu id, after emptying the cache: 583
Out[27]:
[GPUMemory(total=12212, free=11525, used=687),
 GPUMemory(total=12196, free=12185, used=10),
 GPUMemory(total=12196, free=12185, used=10),
 GPUMemory(total=12212, free=11629, used=583)]

Then let's switch up to double the size. So we need to also halve the batch size to avoid running out of GPU memory, and freeze again, and train some more.

Phase 2 a:

lr = 0.001

In [32]:
learn.data = data_phase2
learn.freeze()
gc.collect();
In [33]:
! mkdir ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train/models/
! cp ../../../../../SCRATCH2/marvande/data/train/HR/models/sim9_1b.pth ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train/models/
mkdir: cannot create directory '../../../../../SCRATCH2/marvande/data/train/HR/small-125/train/models/': File exists
In [34]:
learn.load('sim{}_1b'.format(sim));
In [35]:
do_fit('sim{}_2a'.format(sim))
epoch train_loss valid_loss pixel feat_0 feat_1 feat_2 gram_0 gram_1 gram_2 time
0 0.946680 0.996610 0.097366 0.179480 0.164168 0.034117 0.332142 0.175253 0.014084 12:53
1 0.836294 0.907238 0.095648 0.176583 0.160414 0.032235 0.277526 0.151734 0.013097 12:46
2 0.811457 0.881179 0.095342 0.174990 0.159493 0.032045 0.262323 0.144468 0.012518 12:47
3 0.801622 0.845313 0.096339 0.171067 0.156871 0.031763 0.238587 0.138431 0.012254 12:47
4 0.774028 0.818715 0.096593 0.169727 0.155309 0.031179 0.223519 0.130445 0.011943 12:47
5 0.783141 0.855978 0.094315 0.170642 0.156179 0.031227 0.247485 0.143988 0.012142 12:47
6 0.773501 0.840269 0.096518 0.170614 0.155298 0.031074 0.238070 0.136767 0.011929 12:47
7 0.767128 0.841539 0.092560 0.169698 0.154382 0.031009 0.247167 0.134909 0.011813 12:47
8 0.759902 0.801047 0.094216 0.168184 0.153548 0.030862 0.214303 0.128395 0.011538 12:47
9 0.713906 0.776364 0.093549 0.166653 0.152044 0.030681 0.200701 0.121490 0.011246 12:47
In [36]:
shutil.copy(
    path_lr/'models/sim{}_2a.pth'.format(sim),
    '../../../../../SCRATCH2/marvande/data/train/HR/models/')
Out[36]:
'../../../../../SCRATCH2/marvande/data/train/HR/models/sim9_2a.pth'

Phase 2 b:

smaller learning rates, lr = [1e-06 to 0.0001]

In [30]:
learn.load('sim{}_2a'.format(sim));
In [31]:
learn.unfreeze()
In [32]:
do_fit('sim{}_2b'.format(sim), slice(1e-6,1e-4))
epoch train_loss valid_loss pixel feat_0 feat_1 feat_2 gram_0 gram_1 gram_2 time
0 0.899480 0.856622 0.085192 0.138910 0.143318 0.029351 0.255425 0.190546 0.013882 02:50
1 0.801987 0.739572 0.078939 0.132001 0.133066 0.027535 0.195816 0.159238 0.012977 02:43
2 0.702795 0.650789 0.072503 0.125412 0.124186 0.026221 0.153928 0.136299 0.012240 02:43
3 0.645889 0.603309 0.068056 0.119542 0.117616 0.025072 0.136360 0.124983 0.011679 02:43
4 0.612923 0.574370 0.066212 0.115108 0.113188 0.024279 0.125389 0.118911 0.011284 02:44
5 0.590305 0.554575 0.064248 0.111736 0.109828 0.023638 0.119339 0.114800 0.010987 02:43
6 0.575780 0.539744 0.062705 0.108839 0.106975 0.023093 0.115823 0.111553 0.010755 02:43
7 0.567059 0.530553 0.062073 0.107442 0.105702 0.022829 0.111700 0.110230 0.010576 02:43
8 0.560166 0.525502 0.062063 0.106363 0.104559 0.022682 0.110608 0.108720 0.010507 02:43
9 0.557891 0.520142 0.061213 0.105753 0.103958 0.022497 0.108613 0.107703 0.010404 02:43
In [33]:
shutil.copy(
    path_lr/'models/sim{}_2b.pth'.format(sim),
    '../../../../../SCRATCH2/marvande/data/train/HR/models/')
Out[33]:
'../../../../../SCRATCH2/marvande/data/train/HR/models/sim9_2b.pth'

Testing :

In [78]:
# Create target directory & all intermediate directories if don't exists
try:
    os.makedirs('data/')   
    os.makedirs('images/')
    os.makedirs('data/metrics')    
    print("Directory  Created ")
except FileExistsError:
    print("Directory already exists")  

try:
    os.makedirs('data/interpolatedHR')
    os.makedirs('data/interpolatedHR/HR')    
    os.makedirs('data/interpolatedHR/LR')    

    print("Directory  Created ")
except FileExistsError:
    print("Directory  already exists")  

try:
    os.makedirs('data/predictions')  
    os.makedirts('data/predictions/HR')
    os.makedirts('data/predictions/LR')
    print("Directory  Created ")
except FileExistsError:
    print("Directory already exists")    
Directory already exists
Directory  already exists
Directory already exists

Create test data:

Create testing data, data_mr and data_lr both upsampled to size (3, 502, 672), except data_mr is upsampled from size (3, 250, 334) while data_lr is upsampled from size (3, 125, 167). This way, we will test the model on both sets and see if one of them creates better HR images.

In [79]:
size_hr = (3, 502, 672)
bs = 1

#upsampled by factor 4:
print('Upsample factor 4')
print(f'LR test from {path_lr_test}')
print(f'LR train from {path_lr}')
data_lr = get_data(bs, size_hr, path_lr_test, path_lr, path_hr)
print('')
print('Upsample factor 2')
#upsampled by factor 2: 
data_mr = get_data(bs, size_hr, path_mr_test, path_mr, path_hr)
print(f'LR test from {path_mr_test}')
print(f'LR train from {path_mr}')
Upsample factor 4
LR test from ../../../../../SCRATCH2/marvande/data/train/HR/small-125/test
LR train from ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train

Upsample factor 2
LR test from ../../../../../SCRATCH2/marvande/data/train/HR/small-250/test
LR train from ../../../../../SCRATCH2/marvande/data/train/HR/small-250/train
In [80]:
data_mr_test = get_data(bs, (1,3,502,672), path_mr_test, path_mr_test, path_test)
data_lr_test = get_data(bs, (1,3,502,672), path_lr_test, path_lr_test, path_test)
In [81]:
# original HR test data:
data_hr = get_data(bs, size_hr, path_test, path_hr, path_hr)
print(f'LR test from {path_test}')
print(f'LR train from {path_hr}')
LR test from ../../../../../SCRATCH2/marvande/data/train/HR/HR_patches_test/cut_images/tiff_files
LR train from ../../../../../SCRATCH2/marvande/data/train/HR/HR_patches_train/tiff_files

If model was somehow deleted from the path_lr folder in the process somewhere in the notebook, this will copy it from the folder where we saved it to earlier.

In [82]:
# Create target directory & all intermediate directories if don't exists
path = '../../../../../SCRATCH2/marvande/data/train/HR/small-125/train/models/'
try:
    os.makedirs(path)    
    print("Directory " , path ,  " Created ")
except FileExistsError:
    print("Directory " , path ,  " already exists")  

# copy model to right directory to load it:
shutil.copy(
    '../../../../../SCRATCH2/marvande/data/train/HR/models/sim{}_2b.pth'.format(sim),
    path_lr/'models/')
Directory  ../../../../../SCRATCH2/marvande/data/train/HR/small-125/train/models/  already exists
Out[82]:
'../../../../../SCRATCH2/marvande/data/train/HR/small-125/train/models/sim9_2b.pth'

Load the model from the last phase.

In [83]:
learn.load('sim{}_2b'.format(sim))
learn.data = data_mr

Test on one image:

We first evaluate what the model predicts from one MR and LR image (the same obviously).

In [84]:
num_im_test = 10
# Ground truth HR image:
mr = data_mr.test_ds.x.items[num_im_test]
im_mr = open_image(mr)

hr = data_hr.test_ds.x.items[num_im_test]
#transform to tensor:
HR_tens = open_image(hr).data

# LR version of the same image:
lr = data_lr.test_ds.x.items[num_im_test]
im_lr = open_image(lr)

# Check if it's the same image:
pattern_mr_lr = "test\/(.*?)\.tif"
pattern_hr = "tiff_files\/(.*?)\.tif"
assert (re.search(pattern_mr_lr,
                  str(mr)).group(1) == re.search(pattern_mr_lr,
                                                 str(lr)).group(1))
assert (re.search(pattern_hr,
                  str(hr)).group(1) == re.search(pattern_mr_lr,
                                                 str(lr)).group(1))

Upsample test image:

To be able to apply the model and predict the HR image from the LR, we need to resample the LR and MR image to be of the same size as the HR, as the model inputs and outputs images of the same size.

In [85]:
#resize LR to same size as HR:
img_lr = PIL.Image.open(lr).resize((672, 502),
                                   resample=PIL.Image.BICUBIC).convert('RGB')
LR_tens = scaler(trans1(img_lr))
assert (LR_tens.shape == HR_tens.shape)

#resize MR to same size as HR:
img_mr = PIL.Image.open(mr).resize((672, 502),
                                   resample=PIL.Image.BICUBIC).convert('RGB')
MR_tens = scaler(trans1(img_mr))
assert (MR_tens.shape == HR_tens.shape)

print('HR and resized MR and LR shape {}, {}, {}'.format(list(HR_tens.shape), list(MR_tens.shape),
                                              list(LR_tens.shape)))
print('')

#evaluate metrics between ups LR, HR:
evaluate_metrics(HR_tens.numpy(), LR_tens.numpy(), 'ups LR VS HR')
print('')
#evaluate metrics between ups MR, HR:
evaluate_metrics(HR_tens.numpy(), MR_tens.numpy(), 'ups MR VS HR')

plot_triple(MR_tens.numpy()[0, :, :],
            HR_tens.numpy()[0, :, :],
            LR_tens.numpy()[0, :, :], 'ups MR', 'HR', 'ups LR')
HR and resized MR and LR shape [3, 502, 672], [3, 502, 672], [3, 502, 672]

ups LR VS HR: Frobenius diff: 32.5329 frob, Nuclear diff: 390.8117
ups LR VS HR: MAE: 9.860911, MSE: 204.3631, NMSE: 0.0597, PSNR: 25.0268, SSIM: 0.8685

ups MR VS HR: Frobenius diff: 23.4555 frob, Nuclear diff: 300.6457
ups MR VS HR: MAE: 8.156306, MSE: 107.7225, NMSE: 0.0433, PSNR: 27.8077, SSIM: 0.9204

Prediction on test image:

Now that we have resized the LR and MR to the same size as the ground truth HR, we can feed it to the model and predict a new HR image.

In [137]:
# Prediction of model for ups LR:
p_lr, img_pred_lr, b_lr = learn.predict(Image(LR_tens))

# Prediction of model for ups MR:
p_mr, img_pred_mr, b_mr = learn.predict(Image(MR_tens))

# Assert reconstructed HR has same shape as ground truth HR:
assert (list(p_mr.shape) == list(HR_tens.shape))
assert (list(p_mr.shape) == list(HR_tens.shape))

Then we evaluate visually and mathematically how similar the predicted image is to the ground truth HR. We compare those metrics to the similarity metrics computed between the HR images created by simple bicubic interpolation from LR and MR and HR. Ideally, our model should do better than the interpolated images.

In [87]:
#evaluate metrics between pred HR HR from ups LR:
evaluate_metrics(HR_tens.numpy(), LR_tens.numpy(), 'ups LR VS HR')
evaluate_metrics(HR_tens.numpy(),
                 scaler(img_pred_lr).numpy(), 'pred HR from LR VS HR')

print('')
#evaluate metrics between pred HR HR from ups MR:
evaluate_metrics(HR_tens.numpy(), MR_tens.numpy(), 'ups MR VS HR')
evaluate_metrics(HR_tens.numpy(),
                 scaler(img_pred_mr).numpy(), 'pred HR from MR VS HR')



plot_triple(img_pred_lr.numpy()[0, :, :],
            HR_tens.numpy()[0, :, :],
            LR_tens.numpy()[0, :, :], 'pred HR from LR', 'HR', 'ups LR')

plot_triple(img_pred_mr.numpy()[0, :, :],
            HR_tens.numpy()[0, :, :],
            MR_tens.numpy()[0, :, :], 'pred HR from MR', 'HR', 'ups MR')
ups LR VS HR: Frobenius diff: 32.5329 frob, Nuclear diff: 390.8117
ups LR VS HR: MAE: 9.860911, MSE: 204.3631, NMSE: 0.0597, PSNR: 25.0268, SSIM: 0.8685
pred HR from LR VS HR: Frobenius diff: 38.3930 frob, Nuclear diff: 393.1177
pred HR from LR VS HR: MAE: 15.314659, MSE: 286.8343, NMSE: 0.0707, PSNR: 23.5545, SSIM: 0.8706

ups MR VS HR: Frobenius diff: 23.4555 frob, Nuclear diff: 300.6457
ups MR VS HR: MAE: 8.156306, MSE: 107.7225, NMSE: 0.0433, PSNR: 27.8077, SSIM: 0.9204
pred HR from MR VS HR: Frobenius diff: 42.4859 frob, Nuclear diff: 336.1592
pred HR from MR VS HR: MAE: 17.531969, MSE: 348.6843, NMSE: 0.0780, PSNR: 22.7065, SSIM: 0.9139

Visually our outputs look better though our metrics are worse though. This might be to their darker contrasts.

Evaluate on several images:

Now we evaluate the same thing as before but on a bunch of test images in order to get average prediction evaluations.

Predict and interpolate all test images:

Create predictions from MR and LR test images and interpolated HR from MR and LR and save them.

In [88]:
# paths to interpolated and predicted data:
path_int_mr = Path('data/interpolatedHR/MR')
path_pred_mr = Path('data/predictions/MR')
path_int_lr = Path('data/interpolatedHR/LR')
path_pred_lr = Path('data/predictions/LR')

create_pred_int(learn, path_lr_test, path_test, path_int_lr, path_pred_lr)
create_pred_int(learn, path_mr_test, path_test, path_int_mr, path_pred_mr)

Plot different predictions from MR and LR:

In [89]:
## Plot for several:
LR_list = ImageList.from_folder(path_pred_lr).items
MR_list = ImageList.from_folder(path_pred_mr).items
HR_list = ImageList.from_folder(path_test).items

plot_several_tests_2(learn, LR_list, MR_list, path_test)

Traditional loss and error metrics:

Evaluate traditional measures on 20 test images.

Prediction from upsampled LR images (3, 125, 167):
In [90]:
#pick a random number of test images:
NUM_IMAGES = 20
indices = random.sample(
    range(0, len(ImageList.from_folder(path_lr_test).items)), NUM_IMAGES)

df_lr = eval_test_images_df(learn, path_lr_test, path_test, 'LR', indices)
df_lr.to_csv('data/metrics/metrics_sim{}_lr.csv'.format(sim))
Prediction from upsampled MR images (3, 250, 334):
In [91]:
# predict for the same 20 random images:
df_mr = eval_test_images_df(learn, path_mr_test, path_test, 'MR', indices)
df_mr.to_csv('data/metrics/metrics_sim{}_mr.csv'.format(sim))

Compute the average and median of measures if desired:

In [92]:
av_metrics = median_avg_df(df_lr, 'LR').merge(median_avg_df(df_mr, 'MR'),
                                              on='metrics')
av_metrics_ups = median_avg_df(df_lr, 'LR',
                               ups=True).merge(median_avg_df(df_mr,
                                                             'MR',
                                                             ups=True),
                                               on='metrics')
av_metrics = av_metrics.merge(av_metrics_ups, on='metrics')
av_metrics.index = av_metrics['metrics']
av_metrics = av_metrics.drop('metrics', axis=1)
av_metrics.to_csv('data/metrics/av_metrics_sim{}.csv'.format(sim))
av_metrics;

We also plot the different values for each test image for MR and LR to compare to the interpolated values.

In [93]:
plot_metrics(df_lr,'LR', sim)
plot_metrics(df_mr,'MR',sim)

Feature and gram matrix loss:

Evaluate losses used in the network.

  • Feature loss
  • Gram matrix loss
  • Pixel loss
  • Sum of all losses
Prediction from upsampled LR images (3, 250, 334):
In [94]:
# need to recreate the test data from before because we need to 
# resize it to a special size to evaluate the feature loss... 

# create data where x = interpolated HR from LR and y = GT HR
data_int_lr = get_data(20, (1,3,502,672), path_int_lr, path_int_lr, path_test)

# create data where x = predicted HR from LR and y = GT HR
data_pred_lr = get_data(20, (1,3,502,672), path_pred_lr, path_pred_lr, path_test)
In [95]:
# compute data losses between interpolated HR from LR and GT HR
df_int_lr = losses_df(data_int_lr,vgg_m, blocks)

# compute data losses between predicted HR from LR and GT HR
df_pred_lr = losses_df(data_pred_lr,vgg_m, blocks)
In [96]:
plot_metrics_feat(df_pred_lr,df_int_lr,'LR', sim)
Prediction from upsampled MR images (3, 250, 334):
In [97]:
# create data where x = interpolated HR from MR and y = GT HR
data_int_mr = get_data(20, (1,3,502,672), path_int_mr, path_int_mr, path_test)

# create data where x = predicted HR from MR and y = GT HR
data_pred_mr = get_data(20, (1,3,502,672), path_pred_mr, path_pred_mr, path_test)

#data_pred_mr.show_batch(rows = 1,ds_type=DatasetType.Valid)
#data_int_mr.show_batch(rows = 1, ds_type=DatasetType.Valid)
In [98]:
# compute data losses between interpolated HR from MR and GT HR
df_int_mr = losses_df(data_int_mr,vgg_m, blocks)

# compute data losses between predicted HR from MR and GT HR
df_pred_mr = losses_df(data_pred_mr,vgg_m, blocks)
In [99]:
plot_metrics_feat(df_pred_mr,df_int_mr,'MR', sim)
In [100]:
import os
os.system('jupyter nbconvert --to html mod_model_on_superres_norm.ipynb')
Out[100]:
0